package com.example.videopractice.mediaex

import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.media.MediaMuxer
import android.os.Environment
import com.example.base.util.ALog
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.GlobalScope
import kotlinx.coroutines.launch
import java.io.File
import java.util.concurrent.ArrayBlockingQueue

/**
 * H264 编码类
 */
class H264VideoEncoder(private val width: Int,
                       private val height: Int,
                       private val frameRate: Int) {

    private val mediaCodec: MediaCodec
    private val mediaMuxer: MediaMuxer
    private val yuv420Queue = ArrayBlockingQueue<ByteArray>(10)
    private var videoTrack = -1


    var isRunning = false

    init {
        val mediaFormat = MediaFormat.createVideoFormat("video/avc", width, height)
        mediaFormat.setInteger(
            MediaFormat.KEY_COLOR_FORMAT,
            MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar
        )
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, width * height * 5)
        mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, 30)
        mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1)

        mediaCodec = MediaCodec.createEncoderByType("video/avc")
        mediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        mediaCodec.start()

        val path =
            Environment.getExternalStorageDirectory().absolutePath + File.separator + "temp_video.mp4"

        mediaMuxer = MediaMuxer(path, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4) //输出为MP4
    }

    /***
     * 开始编码
     */
    fun startEncoder() {
        GlobalScope.launch(Dispatchers.IO) {
            isRunning = true
            var pts = 0L
            var generateIndex = 0L

            while (isRunning) {
                val input = yuv420Queue.take()
                val yuv420sp = ByteArray(width * height * 3 / 2)
                // 必须要转格式，否则录制的内容播放出来为绿屏
                NV21ToNV12(input, yuv420sp, width, height)

                try {
                    //获取输入流队列
                    val inputBufferIndex = mediaCodec.dequeueInputBuffer(0) //间隔多久 0为立即返回
                    val inputBuffer = mediaCodec.getInputBuffer(inputBufferIndex)
                    if (inputBuffer != null) {
                        pts = computePresentationTime(generateIndex)
                        //输入流入队列
                        inputBuffer.put(yuv420sp)
                        mediaCodec.queueInputBuffer(
                            inputBufferIndex,
                            0,
                            yuv420sp.size,
                            pts,
                            0
                        )
                        generateIndex += 1
                    }

                    //获取输出流队列
                    val bufferInfo = MediaCodec.BufferInfo()
                    var outputBufferIndex =
                        mediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC) //timeoutUs 超时时间

                    if (outputBufferIndex == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                        videoTrack = mediaMuxer.addTrack(mediaCodec.outputFormat)
                        ALog.e("xiao", "format改变, videoTrack: $videoTrack")
                        if (videoTrack >= 0) {
                            mediaMuxer.start()
                            ALog.e("xiao", "开始混合")
                        }
                    }
                    while (outputBufferIndex >= 0) {
                        if (videoTrack >= 0) {
                            mediaCodec.getOutputBuffer(outputBufferIndex)?.let {
                                mediaMuxer.writeSampleData(videoTrack, it, bufferInfo)
                            }
                            mediaCodec.releaseOutputBuffer(outputBufferIndex, false)
                        }
                        outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC)
                    }
                } catch (e: Exception) {
                    ALog.e("xiao", "错误了: ${e.message}")
                    e.printStackTrace()
                }
            }

            try {
                mediaMuxer.stop()
                mediaMuxer.release()
            } catch (e: Exception) {
                e.printStackTrace()
            }

            try {
                mediaCodec.stop()
                mediaCodec.release()
            } catch (e: Exception) {
                e.printStackTrace()
            }

            ALog.e("xiao", "结束")
        }
    }

    /**
     * 停止编码数据
     */
    fun stopEncoder() {
        isRunning = false
    }

    /**
     * 根据帧数生成时间戳
     */
    private fun computePresentationTime(frameIndex: Long): Long {
        return 132 + frameIndex * 1000000 / frameRate
    }

    private fun NV21ToNV12(nv21: ByteArray, nv12: ByteArray, width: Int, height: Int) {
        val frameSize = width * height
        System.arraycopy(nv21, 0, nv12, 0, frameSize)
        for (i in 0 until frameSize) {
            nv12[i] = nv21[i]
        }
        for (j in 0 until (frameSize / 2) step 2) {
            nv12[frameSize + j - 1] = nv21[j + frameSize]
        }
        for (j in 0 until (frameSize / 2) step 2) {
            nv12[frameSize + j] = nv21[j + frameSize - 1]
        }
    }


    fun putData(buffer: ByteArray) {
        if (yuv420Queue.size >= 10) {
            yuv420Queue.poll()
        }
        yuv420Queue.put(buffer)
    }


    companion object {
        private const val TIMEOUT_USEC = 12000L
    }
}